import torch
import numpy as np
from eeg_resnet import Bottleneck, ResNet
 
## 此程序文件对模型训练遍历过程中产生的模型权重文件进行效果测试 ##

device = torch.device('cuda:0' if torch.cuda.is_available() else 'cpu')  # 判断是否有GPU
 
test_data = np.loadtxt('./test_data.txt')  # 测试集的路径
test_labels=np.loadtxt('./test_labels.txt')
 
model = ResNet().to(device)

for k in range(0,50):
    
    model.load_state_dict(torch.load('ResNet_{0}.pth'.format(k)))  # 加载模型
    model.eval()
     
    test_length = test_data.shape[0]

    j=0
     
    for i in range(test_length):
        td = torch.from_numpy(test_data[i]).float().to(device)
        Predict_output = model(td).to(device)
        _, predicted = Predict_output.max(1)
        pred_type = predicted.item()

        if pred_type==test_labels.argmax(1)[i]:
            j=j+1
            
    print('total:',k,':',j/test_length*100,'%')
    k=k+1

